

""" 这里使用Fashion Minist数据集做分类测试 """

import torch
import torchvision

from torch import nn
from torch.utils import data
from torchvision import transforms

import numpy as np
import cv2 as cv # 借助opencv来显示曲线图片。

import TrainAndTest


def show_image(img_tensor, name="img"): 
    """绘制图像,imgs 必须是 三维张量，[通道][行][列]"""

    # 将张量转换为numpy数组
    image_numpy = img_tensor.detach().cpu().numpy()
    # OpenCV 默认的图像数据格式是 (H, W, C)，而深度学习框架中的张量图像是 (C, H, W)
    image_numpy = np.transpose(image_numpy, (1, 2, 0))
    cv.imshow(name, cv.resize(image_numpy,(280,280)))
    cv.waitKey(0)
    cv.destroyAllWindows()



if __name__ == "__main__":
    # 下载/加载数据文件。
    trans = transforms.ToTensor()
    #mnist_train = torchvision.datasets.FashionMNIST(
    #    root="../data", train=True, transform=trans, download=True)
    mnist_test = torchvision.datasets.FashionMNIST(
        root="../data", train=False, transform=trans, download=True)
    #print( "mnist_train = ",len(mnist_train)," , mnist_test = ", len(mnist_test))
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
 
    # 一次加载所有的数据集合。
    #D = data.DataLoader(mnist_test, batch_size=len(mnist_test),num_workers=1)
    D = data.DataLoader(mnist_test, batch_size=len(mnist_test),num_workers=1)
    iD = iter(D)
    # ! pip install d2l 的numpy在windows运行会有个错误： next(iter(data.DataLoader))  行： TypeError: expected np.ndarray (got numpy.ndarray) 。
    # ! 需要用conda重新安装numpy 1.21.5 。
    
    # for X,y in iD:
    #     print(X.shape)
    #     print(y.shape)
       
    X, y = next(iD)
    print(X.shape)
    print(y.shape)
    # 显示一张图片：
    #show_image(X[0,:,:,:],text_labels[y[0]])


    # 下面先测试 get_data_iter
    sample_list = []
    lable_list = []
    batch_size = 256
    for i in range(min(len(X),2000)):
        sample_list.append(X[i,:,:,:])
        lable_list.append(y[i])
    train_iter = TrainAndTest.get_data_iter(sample_list,lable_list,batch_size)

    sample_list = []
    lable_list = []
    batch_size = 256
    for i in range(min(len(X),2000)):
        sample_list.append(X[len(X)-i-1,:,:,:])
        lable_list.append(y[len(X)-i-1])
    test_iter = TrainAndTest.get_data_iter(sample_list,lable_list,batch_size)
  
    
    net = nn.Sequential(
        nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.ReLU(),
        nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Conv2d(6, 16, kernel_size=5), nn.ReLU(),
        nn.AvgPool2d(kernel_size=2, stride=2),
        nn.Flatten(),
        nn.Linear(16 * 5 * 5, 10), nn.ReLU()   # 原来的式子，Flatten之后刚好是 16 * 5 * 5 的维度。
       # nn.Linear(16 * 5 * 3, 10), nn.ReLU()   # 修改成这样，就会报错
    )
    # 上面是进一步简化的LeNet
    
    TrainAndTest.train_and_test_Classify_net(net,train_iter,test_iter,10,0.1)



















